from pydantic import BaseModel, Field
from langchain.agents import tool

from config import web_json


class summary_input(BaseModel):
    title: str = Field(
        ...,
        description="从用户输入内容提取的标题",
    )
    desc: str = Field(
        ...,
        description="从用户输入内容提取的描述",
    )


summary_desc = """\
提取适合生成标题和描述的文本段落
"""


@tool(description=summary_desc, args_schema=summary_input)
def summary(*args, **kwargs) -> str:
    title = ""
    if "title" in kwargs:
        title = kwargs["title"]
    else:
        return "输入参数错误，根据工具参数要求输入title参数"

    desc = ""
    if "desc" in kwargs:
        desc = kwargs["desc"]
    else:
        return "输入参数错误，根据工具参数要求输入desc参数"

    res = {
        "title": title,
        "type": "header",
        "data": desc,
    }

    web_json.append(res)

    # return json.dumps(res, ensure_ascii=False)
    return "报告头部生成成功"
